import numpy as np
from typing import List, Tuple, Optional


class SwitchPuzzleSolver:
    def __init__(self, grid: List[List[int]]):
        """
        初始化开关谜题求解器

        Args:
            grid: 3x3的二维数组，0表示无开关，1-关闭，2-等待，3-开启
        """
        self.grid = np.array(grid)
        self.rows = len(grid)
        self.cols = len(grid[0]) if grid else 0
        self.mod = 4  # 模4运算

    def get_adjacent_positions(self, row: int, col: int) -> List[Tuple[int, int]]:
        """获取指定位置的十字相邻位置"""
        adjacent = []
        directions = [(-1, 0), (1, 0), (0, -1), (0, 1)]  # 上下左右

        for dr, dc in directions:
            new_row, new_col = row + dr, col + dc
            if 0 <= new_row < self.rows and 0 <= new_col < self.cols:
                adjacent.append((new_row, new_col))

        return adjacent

    def create_equation_matrix(self) -> Tuple[np.ndarray, np.ndarray]:
        """创建线性方程组的系数矩阵和常数向量"""
        n = self.rows * self.cols
        A = np.zeros((n, n), dtype=int)
        b = np.zeros(n, dtype=int)

        # 为每个位置创建方程
        for row in range(self.rows):
            for col in range(self.cols):
                if self.grid[row, col] == 0:  # 无开关位置
                    continue

                pos_idx = row * self.cols + col

                # 目标状态：从当前状态到开启状态(3)需要的改变次数
                current_state = self.grid[row, col]
                target_state = 3
                required_change = (target_state - current_state) % self.mod
                b[pos_idx] = required_change

                # 当前位置的开关会影响自己和相邻位置
                adjacent_positions = self.get_adjacent_positions(row, col)
                adjacent_positions.append((row, col))  # 包括自己

                for adj_row, adj_col in adjacent_positions:
                    if self.grid[adj_row, adj_col] != 0:  # 相邻位置有开关
                        adj_idx = adj_row * self.cols + adj_col
                        A[pos_idx, adj_idx] = 1

        return A, b

    def gaussian_elimination_mod4(
        self, A: np.ndarray, b: np.ndarray
    ) -> Optional[np.ndarray]:
        """在模4域中进行高斯消元"""
        n = len(A)
        augmented = np.column_stack([A, b])

        # 前向消元
        for i in range(n):
            # 寻找主元
            pivot_row = i
            for j in range(i + 1, n):
                if abs(augmented[j, i]) > abs(augmented[pivot_row, i]):
                    pivot_row = j

            if pivot_row != i:
                augmented[i], augmented[pivot_row] = (
                    augmented[pivot_row].copy(),
                    augmented[i].copy(),
                )

            # 如果主元为0，尝试找到非零主元
            if augmented[i, i] == 0:
                for j in range(i + 1, n):
                    if augmented[j, i] != 0:
                        augmented[i], augmented[j] = (
                            augmented[j].copy(),
                            augmented[i].copy(),
                        )
                        break
                else:
                    # 检查是否有解
                    if augmented[i, -1] != 0:
                        return None  # 无解
                    continue

            # 消元
            for j in range(i + 1, n):
                if augmented[j, i] != 0:
                    pivot_val = int(augmented[i, i])
                    current_val = int(augmented[j, i])

                    # 处理模4中的特殊情况
                    if pivot_val == 1:
                        factor = current_val
                    elif pivot_val == 3:
                        factor = (current_val * 3) % 4
                    elif pivot_val == 2:
                        # 2没有逆元，需要特殊处理
                        if current_val == 1:
                            factor = 2  # 2 * 2 = 0 (mod 4), 所以用2
                        elif current_val == 3:
                            factor = 2  # 2 * 2 = 0 (mod 4), 所以用2
                        else:
                            factor = 0
                    else:
                        factor = 0

                    augmented[j] = (augmented[j] - factor * augmented[i]) % self.mod

        # 回代求解
        x = np.zeros(n, dtype=int)
        for i in range(n - 1, -1, -1):
            if augmented[i, i] == 0:
                if augmented[i, -1] != 0:
                    return None  # 无解
                x[i] = 0  # 自由变量设为0
                continue

            sum_val = 0
            for j in range(i + 1, n):
                sum_val = (sum_val + augmented[i, j] * x[j]) % self.mod

            pivot_val = int(augmented[i, i])
            const_val = int(augmented[i, -1])

            # 处理模4中的特殊情况
            if pivot_val == 1:
                x[i] = (const_val - sum_val) % self.mod
            elif pivot_val == 3:
                x[i] = ((const_val - sum_val) * 3) % self.mod
            elif pivot_val == 2:
                # 2没有逆元，检查是否有解
                target = (const_val - sum_val) % self.mod
                if target % 2 == 0:
                    x[i] = target // 2  # 如果目标值是偶数，可以除以2
                else:
                    return None  # 无解
            else:
                x[i] = 0

        return x

    def solve(self) -> Optional[List[Tuple[int, int]]]:
        """求解开关谜题"""
        A, b = self.create_equation_matrix()
        solution = self.gaussian_elimination_mod4(A, b)

        if solution is None:
            return None

        # 将解转换为开关位置
        moves = []
        for i, count in enumerate(solution):
            if count > 0:  # 需要按动
                row = i // self.cols
                col = i % self.cols
                for _ in range(count):
                    moves.append((row, col))

        return moves

    def simulate_moves(self, moves: List[Tuple[int, int]]) -> np.ndarray:
        """模拟执行移动后的最终状态"""
        final_grid = self.grid.copy()

        for row, col in moves:
            if final_grid[row, col] == 0:
                continue

            # 改变当前位置：1→2→3→1
            if final_grid[row, col] == 1:
                final_grid[row, col] = 2
            elif final_grid[row, col] == 2:
                final_grid[row, col] = 3
            elif final_grid[row, col] == 3:
                final_grid[row, col] = 1

            # 改变相邻位置
            adjacent_positions = self.get_adjacent_positions(row, col)
            for adj_row, adj_col in adjacent_positions:
                if final_grid[adj_row, adj_col] != 0:
                    if final_grid[adj_row, adj_col] == 1:
                        final_grid[adj_row, adj_col] = 2
                    elif final_grid[adj_row, adj_col] == 2:
                        final_grid[adj_row, adj_col] = 3
                    elif final_grid[adj_row, adj_col] == 3:
                        final_grid[adj_row, adj_col] = 1

        return final_grid


def print_grid(grid: np.ndarray, title: str = ""):
    """打印网格状态（支持任意大小）"""
    if title:
        print(f"\n{title}:")

    rows, cols = grid.shape

    # 打印顶部边框
    print("┌" + "─────┬" * (cols - 1) + "─────┐")

    for i in range(rows):
        print("│", end=" ")
        for j in range(cols):
            if grid[i, j] == 0:
                print("   ", end=" │ ")
            else:
                print(f" {grid[i, j]} ", end=" │ ")
        print()

        # 打印分隔线（除了最后一行）
        if i < rows - 1:
            print("├" + "─────┼" * (cols - 1) + "─────┤")

    # 打印底部边框
    print("└" + "─────┴" * (cols - 1) + "─────┘")


def get_user_input() -> List[List[int]]:
    """获取用户输入的九宫格状态"""
    print("请输入九宫格状态（3x3矩阵）：")
    print("状态说明: 0=无开关, 1=关闭, 2=等待, 3=开启")
    print("请输入每行，用空格分隔数字，例如：1 0 1")
    print()

    grid = []
    for i in range(3):
        while True:
            try:
                row_input = input(f"第{i+1}行: ").strip()
                row_values = [int(x) for x in row_input.split()]

                if len(row_values) != 3:
                    print("错误：每行必须输入3个数字，请重新输入")
                    continue

                # 验证输入值
                for val in row_values:
                    if val not in [0, 1, 2, 3]:
                        print("错误：状态值只能是 0, 1, 2, 3，请重新输入")
                        break
                else:
                    grid.append(row_values)
                    break

            except ValueError:
                print("错误：请输入有效的数字，用空格分隔")

    return grid


def main():
    """主函数"""
    print("三相开关谜题求解器")
    print("=" * 30)
    print("目标：将所有开关变为开启状态(3)")
    print()

    # 获取用户输入
    grid = get_user_input()

    print("\n您输入的九宫格状态：")
    print_grid(np.array(grid))

    # 求解
    solver = SwitchPuzzleSolver(grid)
    moves = solver.solve()

    if moves:
        print(f"\n✅ 找到解决方案！需要按动 {len(moves)} 次")
        print("\n具体操作步骤：")

        # 按步骤显示操作
        current_grid = grid.copy()
        for i, (row, col) in enumerate(moves, 1):
            print(f"\n第{i}步：按动位置 ({row+1}, {col+1})")

            # 模拟这一步的效果
            if current_grid[row][col] != 0:
                # 改变当前位置：1→2→3→1
                if current_grid[row][col] == 1:
                    current_grid[row][col] = 2
                elif current_grid[row][col] == 2:
                    current_grid[row][col] = 3
                elif current_grid[row][col] == 3:
                    current_grid[row][col] = 1

                # 改变相邻位置
                adjacent_positions = solver.get_adjacent_positions(row, col)
                for adj_row, adj_col in adjacent_positions:
                    if current_grid[adj_row][adj_col] != 0:
                        if current_grid[adj_row][adj_col] == 1:
                            current_grid[adj_row][adj_col] = 2
                        elif current_grid[adj_row][adj_col] == 2:
                            current_grid[adj_row][adj_col] = 3
                        elif current_grid[adj_row][adj_col] == 3:
                            current_grid[adj_row][adj_col] = 1

            print_grid(np.array(current_grid), f"第{i}步后的状态")

        # 验证最终结果
        final_grid = solver.simulate_moves(moves)
        print_grid(final_grid, "最终状态")

        # 检查是否所有开关都开启
        all_on = True
        for i in range(solver.rows):
            for j in range(solver.cols):
                if final_grid[i, j] != 0 and final_grid[i, j] != 3:
                    all_on = False
                    break

        if all_on:
            print("\n🎉 成功！所有开关都已开启！")
        else:
            print("\n⚠️  警告：最终状态不正确，可能存在计算错误")

    else:
        print("\n❌ 无解！无法将所有开关变为开启状态")
        print("提示：某些初始状态可能无法达到目标状态")


if __name__ == "__main__":
    main()
